import sys
import os
import pylab as plt
from params import *
import math

edges = set()
PA = 1
GM = 2
ORIG = 3
def readInitialEdges():
  filename = os.path.join(WORKDIR, 'gowalla', 'temporal', 'data', 'gowalla_temporal_graph.txt')
  for line in open(filename):
    n1, n2, day = map(int, line.strip().split())
    if day > 0:
      break
    if n1 > n2:
      n1, n2 = n2, n1
    edges.add((n1,n2))

def IsInitialEdge(n1,n2):
  if n1 > n2:
    n1, n2 = n2, n1
  return (n1,n2) in edges

def scale_distance(d):
    M1 = 10**(-1.0)
    M2 = 10**(4.4)
    N = 1000
    d = max(d, M1)
    d = min(d, M2)
    i = (d-M1)/(M2-M1)
    index = int(N*i)
    i = float(index)/N
    v = M1 + i*(M2-M1)
    return v

def scale_distance_log(d):
    M1 = -1.0
    M2 = 4.4
    N = 100
    d = max(d, 10**M1)
    d = min(d, 10**M2)
    l = math.log10(d)
    i = (l-M1)/(M2-M1)
    index = int(N*i)
    i = float(index)/N
    l2 = M1 + i*(M2-M1)
    v = 10**l2
    return v

def GetDegreeCorrelation(file):
    data = {}
    def update_user(user, dist):
        if not user in data:
            data[user] = (0.0, 0.0)
        sum_dist, sum_friends = data[user]
        sum_dist += dist
        sum_friends += 1
        data[user] = (sum_dist, sum_friends)

    for line in open(file):
        user1,user2, dist = line.strip().split(' ')
        user1 = int(user1)
        user2 = int(user2)
        if IsInitialEdge(user1, user2):
          continue
        dist = float(dist)
        update_user(user1,dist)
        update_user(user2,dist)

    degrees = {}
    for user in data:
        sum_dist, deg = data[user]
        if not deg in degrees:
            degrees[deg] = []
        degrees[deg].append(sum_dist/deg)
    x =  sorted(degrees)

    def avg(seq):
      return float(sum(seq))/len(seq)

    def geom_avg(seq):
      seq = [max(i,1) for i in seq]
      return math.exp(avg([math.log(i) for i in seq]))

    y = [geom_avg(degrees[k]) for k in x]
    #y = [avg(degrees[k]) for k in x]

    return x,y

def GetDistanceDist(file):
    distances = {}
    for line in open(file):
        user1,user2, dist = line.strip().split(' ')
        user1 = int(user1)
        user2 = int(user2)
        if IsInitialEdge(user1, user2):
          continue
        dist = float(dist)
        dist = scale_distance_log(dist)
        distances.setdefault(dist,0)
        distances[dist] += 1

    x = sorted(distances)
    y = [distances[k] for k in x]

    c = [0.0]
    for i in y:
        c.append(c[-1]+i)

    tot = c[-1]
    c = [1.0*i/tot for i in c[:-1]]
    return x,c

def GetDegreeDist(file):
    data = {}
    def update_user(user):
        if not user in data:
            data[user] = 0
        data[user] += 1

    lines = 0
    for line in open(file):
        lines += 1
        user1,user2 = map(int,line.split(' ')[:2])
        if IsInitialEdge(user1, user2):
          continue
        update_user(user1)
        update_user(user2)

    degrees = {}
    edges = 0
    for user in data:
        deg = data[user]
        edges += deg
        if not deg in degrees:
            degrees[deg] = 0
        degrees[deg] += 1

    x = sorted(degrees)
    y = [degrees[k] for k in x]
    tot = sum(y)
    N = len(data)
    K = edges//2
    avg_deg = 2.0*K/N
    print "Nodes %d, edges %d, avg degree %f"%(N,K,avg_deg)
    y = [float(i)/tot for i in y]
    c = [y[0]]
    for i in y[1:]:
        c.append(c[-1]+i)
    c = [1.0-i for i in c]
    return x,y

#def GetCusteringDist(file):
#    import networkx as nx
#    from math import log
#    G = nx.Graph();
#    
#    lines = 0
#    for line in open(file):
#        lines += 1
#        user1,user2 = map(int,line.split(' ')[:2])
#        if IsInitialEdge(user1, user2):
#            continue
#        G.add_edge(user1,user2)
#
#    coefficients = {}
#    numOfNodes = {}
#    for user in G:
#        coef = nx.clustering(G, user)
#        deg = G.degree(user)
#        #deg = int(log(deg,2))
#        #deg = 2 ** deg
#        if not deg in coefficients:
#            coefficients[deg] = 0.
#            numOfNodes[deg] = 0.
#        coefficients[deg] += coef
#        numOfNodes[deg] += 1.
#    
#    total = 0.;
#    totalNumOfNodes = 0.;
#    for degree in coefficients:
#        total += float(coefficients[degree])
#        totalNumOfNodes += float(numOfNodes[degree])
#        coefficients[degree] = float(coefficients[degree]) / float(numOfNodes[degree])
#    
#    avg_coef =  total / totalNumOfNodes
#    print 'Average Coefficient (without old edges) ' + str(avg_coef)
#    
#    x = sorted(coefficients)
#    y = [coefficients[k] for k in x]
#    return x,y
    
#def GetTriangleDistanceDistr(file):
#    import networkx as nx
#    from math import log
#    G = nx.Graph();
#    
#    lines = 0
#    for line in open(file):
#        lines += 1
#        user1,user2 = map(int,line.split(' ')[:2])
#        dist = float(line.split(' ')[2])
#        if IsInitialEdge(user1, user2):
#            continue
#        G.add_edge(user1,user2, distance=dist)
#
#    triangleDists = {}
#    for user in G:
#        distance_sum = 0
#        num = 0
#        for friend in G[user]:
#          for friend2 in G[friend]:
#            if friend2 in G[user]:
#              num += 1
#              distance = G[user][friend]["distance"]
#              distance += G[friend][friend2]["distance"]
#              distance += G[friend2][user]["distance"]
#              if (distance < 1):
#                distance = 1
#              distance_sum += log(distance)
#        # Append the geom. average of the distances
#        if num > 0:
#          value = int( math.exp(float(distance_sum) / float(num) ))
#          value = 1.2 ** (int(log(value,1.2)))
#          if value in triangleDists:
#            triangleDists[value] += 1.
#          else:
#            triangleDists[value] = 1.
#    
#    total_triangles = sum([triangleDists[value] for value in triangleDists])
#    
#    x = sorted(triangleDists)
#    y = [ float(triangleDists[k]) / float(total_triangles) for k in x]
#    for i in xrange(1,len(y)):
#      y[i]+=y[i-1]
#    return x,y

def GetTriangleDistanceDegreeCorrelation(file):
    import networkx as nx
    from math import log
    G = nx.Graph();
    
    lines = 0
    for line in open(file):
        lines += 1
        user1,user2 = map(int,line.split(' ')[:2])
        dist = float(line.split(' ')[2])
        if IsInitialEdge(user1, user2):
            continue
        G.add_edge(user1,user2, distance=dist)

    def __get_dist(d):
      #return log(max(1.0, d))
      return max(1.0, d)

    def __get_triangle_len(u1, u2, u3):
        distance = __get_dist(G[user][friend]["distance"])
        distance += __get_dist(G[friend][friend2]["distance"])
        distance += __get_dist(G[friend2][user]["distance"])
        distance = distance / 3.0
        return distance

    userTriangleDist = {}
    userTriangleCount = {}
    for user in G:
        for friend in G[user]:
          for friend2 in G[friend]:
            if friend2 in G[user]:
              tri_len = __get_triangle_len(user, friend, friend2)
              userTriangleDist[user] = userTriangleDist.get(user, 0) + tri_len
              userTriangleCount[user] = userTriangleCount.get(user, 0) + 1.0

    triangleDists = {}
    triangleCount = {}
    for user in userTriangleDist:
      degree = G.degree(user)
      num = userTriangleDist[user]
      den = userTriangleCount[user]
      tri_len = num / den
      triangleDists[degree] = triangleDists.get(degree, 0) + log(tri_len)
      triangleCount[degree] = triangleCount.get(degree, 0) + 1.

    x = sorted(triangleDists)
    y = [ math.exp(float(triangleDists[k] /  triangleCount[k])) for k in x]
    return x,y

def GetFileType(filename):
  print 'get type ', filename
#  if not '3' in filename:
#    return None
#  if filename.startswith('pa_graph'):
#    return PA
  if filename.startswith('gravity_graph'):
    return GM
  if filename.startswith('graph_pa_social'):
    return PA
  return None

def GetGraphs(filepath):
  results = {}
  x,c = GetDegreeDist(filepath)
  results['DegreeCDF'] = (x,c)
  x,c = GetDistanceDist(filepath)
  results['DistanceCDF'] = (x,c)
  x,c = GetDegreeCorrelation(filepath)
  results['DegreeCorr'] = (x,c)

  #x,c = GetCusteringDist(filepath)
  #results['DegreeClusteringCorr'] = (x,c)
  #x,c = GetTriangleDistanceDistr(filepath)
  #results['TriangleLengthDist'] = (x,c)
  x,c = GetTriangleDistanceDegreeCorrelation(filepath)
  results['TriangleLengthDegreeCorr'] = (x,c)

  return results

def avg_plot(plot_list):
  bins = {}
  for x,y in plot_list:
    for xv,yv in zip(x,y):
      bins.setdefault(xv,[]).append(yv)

  x = sorted(bins)
  y = [float(sum(bins[k]))/len(bins[k]) for k in x]
  y2 = [float(sum(i**2 for i in bins[k]))/len(bins[k]) for k in x]
  std = [a - b * b for a,b in zip(y2, y)]

  return x,y,std

def fig_plot(plots, graph_type, xlbl, ylbl, filename, loc, xmin=None,
    xmax=None,ymin=None,ymax=None,log=True):
  lbls = ['Data', 'GM', 'PA']
  MSIZE = 6
  MSIZE1 = 6
  MSIZE2 = 4
  plt.figure()
  plt.clf()
  ax = plt.axes(FIG_AXES2)
  #print len(plots[GM][graph_type])
  x1, c1 = plots[ORIG][graph_type]
  x2, c2, e2 = avg_plot(plots[GM][graph_type])
  x3, c3, e3 = avg_plot(plots[PA][graph_type])
  if log:
    fnc = plt.loglog
  else:
    fnc = plt.semilogx
  fnc(x1,c1,'k.', mfc='None', ms=MSIZE1)
  fnc(x2,c2,'rv', mfc='None', mec='r', ms=MSIZE2)
  fnc(x3,c3,'b+', mfc='None', mec='b', ms=MSIZE2)

  if xmin:
    plt.xlim(xmin=xmin)
  if xmax:
    plt.xlim(xmax=xmax)
  if ymin:
    plt.ylim(ymin=ymin)
  if ymax:
    plt.ylim(ymax=ymax)

  plt.legend(lbls, numpoints=1, loc=loc)

#  low_std2 = [a - 1*b for a,b in zip(c2, e2)]
#  high_std2 = [a + 1*b for a,b in zip(c2, e2)]
#  fnc(x2,low_std2,'r_', mfc='None', ms=MSIZE1)
#  fnc(x2,high_std2,'r_', mfc='None', ms=MSIZE1)
#
#  low_std3 = [a - 1*b for a,b in zip(c3, e3)]
#  high_std3 = [a + 1*b for a,b in zip(c3, e3)]
#  fnc(x3,low_std3,'b_', mfc='None', ms=MSIZE1)
#  fnc(x3,high_std3,'b_', mfc='None', ms=MSIZE1)

  plt.grid(True)
  plt.xlabel(xlbl)
  plt.ylabel(ylbl)
  plt.savefig(filename)
  plt.close()

def main():
  readInitialEdges()

  DIR = os.path.join(WORKDIR, 'gowalla', 'temporal', 'model_graphs')
  plots = {}
  plots[GM] = {}
  plots[PA] = {}
  for filename in os.listdir(DIR):
    model = GetFileType(filename)
    if not model:
      continue
    print ''
    print model, filename
    #if 'DegreeCDF' in plots[model] and len(plots[model]['DegreeCDF']) > 50:
    #  continue
    filepath = os.path.join(DIR, filename)
    graphs = GetGraphs(filepath)
    for graph_type in graphs:
      print graph_type
      x,y = graphs[graph_type]
      plots[model].setdefault(graph_type, []).append((x,y))

  orig_file = os.path.join(WORKDIR, 'gowalla', 'temporal', 'original_graph.txt')
  orig_graphs = GetGraphs(orig_file)
  plots[ORIG] = orig_graphs

  fig_plot(plots, 'DegreeCDF', 'Node degree', 'PMF',
      'model_degree.pdf',loc='upper right',
      xmin=1, xmax=1e3, ymin = 1e-5)

  fig_plot(plots, 'DistanceCDF', 'Link length [km]', 'CDF',
      'model_distance.pdf',loc='upper left',
      xmin=1, xmax=20000, log=False)

  fig_plot(plots, 'DegreeCorr', 'Node degree', 'Friend distance [km]',
      'model_correlation.pdf', loc = 'lower right', 
      xmin=1, xmax=1e3, ymin=1e2, ymax=1e4)

  fig_plot(plots, 'TriangleLengthDegreeCorr', 'Degree', 'Geo. triangle length [km]',
      'model_triang_length_corr.pdf',loc='lower right',
      xmin=1, xmax=1e3, ymin=1e2, ymax=1e4)

  #fig_plot(plots, 'DegreeClusteringCorr', 'Node degree', 'Avg. clustering coeff.',
  #    'model_clustering_coeff.pdf', loc = 'lower left', xmin=1, xmax=1e4, ymax = 1, ymin = 1e-5)

  #fig_plot(plots, 'TriangleLengthDist', 'Average Triangle Length (km) per User', 'CDF',
  #    'model_triang_length.pdf',loc='lower right', log=False,
  #    xmin=1, xmax=1e4)#, ymin = 0.9*1e-5)
if __name__ == '__main__':
  main()
